#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test results')
parser.add_argument('-o', '--output', help='Output figure')
args = parser.parse_args()

f = addfile('bw.720x360x30.h0.nc')

proj = ccrs.PlateCarree(central_longitude=180)
fig = plt.figure(figsize=(10, 5))

ax = fig.add_subplot(2, 2, 1, projection=proj)
ps = f.phs.isel(time=4) / 100
levels = np.linspace(992, 1006, 8)
cmap = ListedColormap([
  [0.84, 0.25, 0.54],
  [0.15, 0.25, 0.58],
  [0.30, 0.66, 0.35],
  [0.80, 0.86, 0.32],
  [0.98, 0.95, 0.31],
  [0.92, 0.64, 0.25],
  [0.85, 0.23, 0.19],
])
plot_contour_map(ax, ps, cmap=cmap, levels=levels, left_string='PS at Day 4')

ax = fig.add_subplot(2, 2, 2, projection=proj)
ps = f.phs.isel(time=6) / 100
plot_contour_map(ax, ps, cmap=cmap, levels=levels, left_string='PS at Day 6')

ax = fig.add_subplot(2, 2, 3, projection=proj)
ps = f.phs.isel(time=8) / 100
levels = np.linspace(930, 1020, 11)
cmap = ListedColormap([
  [0.84, 0.25, 0.54],
  [0.60, 0.18, 0.54],
  [0.15, 0.25, 0.58],
  [0.23, 0.51, 0.77],
  [0.30, 0.67, 0.91],
  [0.30, 0.66, 0.35],
  [0.80, 0.85, 0.32],
  [0.98, 0.95, 0.31],
  [0.92, 0.64, 0.25],
  [0.85, 0.23, 0.19],
])
plot_contour_map(ax, ps, cmap=cmap, levels=levels, left_string='PS at Day 8')

ax = fig.add_subplot(2, 2, 4, projection=proj)
ps = f.phs.isel(time=10) / 100
plot_contour_map(ax, ps, cmap=cmap, levels=levels, left_string='PS at Day 10')

plt.tight_layout()
if args.output:
	plt.savefig(args.output)
else:
	plt.show()
